#include <stdio.h>
#include <math.h>

int tileWidth;

//Globals
float *A_GPU, *B_GPU, *C_GPU; //GPU pointers

dim3 dimBlock, dimGrid; //Block and Grid Dimensions

__global__ void MatrixMultiplySharedMemory(float* A, float* B, float* C, int ay, int by)
{
	int tileWidth = blockDim.x;
	
	//extern __shared__ float sharedArrays[];
	//float* Ashared = (float*) &sharedArrays;
	//float* Bshared = (float*) &Ashared[tileWidth*tileWidth];

	__shared__ float Ashared[4];
	__shared__ float Bshared[4];

	int tx = threadIdx.x;
	int ty = threadIdx.y;

	int row = (tileWidth * blockIdx.y) + ty;
	int col = (tileWidth * blockIdx.x) + tx;
	float cVal = 0;
	
	
	for (int i = 0; i < ceil((float)ay / tileWidth); i++)
	{
		Ashared[ty * tileWidth + tx] = A[row * ay + (i * tileWidth + tx)];
		Bshared[ty * tileWidth + tx] = B[(ty + (i * tileWidth)) * ay + col];

		__syncthreads();

		for(int j = 0; j < tileWidth; j++)
		{
			cVal += Ashared[ty * tileWidth + j] * Bshared[j * tileWidth + tx];
		}
		
		__syncthreads();
	}
	__syncthreads();
	C[row * ay + col] = cVal;
	
}

int main()
{

	cudaError_t err = cudaSuccess;		
	err = cudaSetDevice(1);	
	if (err == cudaSuccess) { printf("\nNo problems setting to device 1!\n");}

	tileWidth = 2;	
	dimBlock.x = tileWidth;
	dimBlock.y = tileWidth;
	dimBlock.z = 1;

	int ax = 4;	
	int ay = 4;

	int bx = 4;
	int by = 2;
	
	int cx = ax;
	int cy = by;



	float* A = (float *)malloc(ax*ay*sizeof(float));
	float* B = (float *)malloc(bx*by*sizeof(float));
	float* C = (float *)malloc(cx*cy*sizeof(float));	

	int i,j;
	
	// Setup padded matrices
	int axDim = ax;
	int ayDim = ay;
	int bxDim = bx;
	int byDim = by;
	int cxDim = ax;
	int cyDim = by;
	
	if (ax%2 != 0) { cxDim = axDim = ax - (ax%2) + 2; }
	if (ay%2 != 0) { ayDim = ay - (ay%2) + 2; }
	if (bx%2 != 0) { bxDim = bx - (bx%2) + 2; }
	if (by%2 != 0) { cyDim = byDim = by - (by%2) + 2; }
	
	dimGrid.x = cxDim;//ceil((float)(ayDim)/tileWidth);
	dimGrid.y = cyDim;//ceil((float)(ayDim)/tileWidth);

	float* _A = (float *)malloc(axDim*ayDim*sizeof(float));
	float* _B = (float *)malloc(bxDim*byDim*sizeof(float));
	float* _C = (float *)malloc(cxDim*cyDim*sizeof(float));

	err = cudaMalloc(&A_GPU, axDim*ayDim*sizeof(float));
	err = cudaMalloc(&B_GPU, bxDim*byDim*sizeof(float));
	err = cudaMalloc(&C_GPU, cxDim*cyDim*sizeof(float));

	if (err == cudaSuccess) { printf("\nNo problems allocating device memory!\n");}

	for(i = 0; i < ax; i++)
	{
		for(j = 0; j < ay; j++)
		{
			A[(ay*i)+j] = (float)2;

			//printf("j= %d ",j);
			//printf(" A[%d] = %f ",ay*i+j,A[(ay*i)+j]);
		}
	}

	for(i = 0; i < bx; i++)
	{
		for(j = 0; j < by; j++)
		{
			B[(by*i)+j] = (float)2;
			//printf("j= %d ",j);
			//printf(" B[%d] = %f ",by*i+j,B[(by*i)+j]);
		}
	}

	for(i = 0; i < cx; i++)
	{
		for(j = 0; j < cy; j++)
		{
			C[(cy*i)+j] = (float)0;
		}	
	}

	for(i = 0; i < axDim; i++)
	{
		for(j = 0; j < ayDim; j++)
		{
			if (i < ax && j < ay)
			{
				_A[(ayDim*i)+j] = A[ay*i+j];
				//_A[i][j]=A[i][j];
				//printf("j= %d ",j);
				//printf(" _A[%d] = A[%d] = %f ",ayDim*i+j, ay*i+j,_A[(ayDim*i)+j]);
			}
			else
			{
				_A[ayDim*i+j] = (float)0;
				//_A[i][j]=(float)0;
			}
			//printf("\n");
		}
	}

	for(i = 0; i < bxDim; i++)
	{
		for(j = 0; j < byDim; j++)
		{
			if (i < bx && j < by)
				{
					_B[byDim*i+j] = B[by*i+j];
					//printf("\nj= %d ",j);
					//printf(" _B[%d] = %f != %f = B[%d] ", byDim*i+j, 
					// _B[(byDim*i)+j], B[by*i+j], by*i+j);
				}
				else
				{
					_B[byDim*i+j] = (float)0;
					//_B[i][j]=(float)0;
				}
			//printf("\n");
			
		}
	}

	for(i = 0; i < cxDim; i++)
	{
		//printf("\n\ni = %d ", i);
		for(j = 0; j < cyDim; j++)
		{
			if (i < cx && j < cy)
			{
				_C[cyDim*i+j] = (float)0;
				//printf("\nj = %d ", j);
				//printf(" _C[%d] = %f ", cyDim*i+j, _C[cyDim*i+j]);
			}
			else
			{
				_C[cyDim*i+j] = (float)0;
			}
			
		}
	}

	err = cudaMemcpy(A_GPU, _A, axDim*ayDim*sizeof(float), cudaMemcpyHostToDevice);
	err = cudaMemcpy(B_GPU, _B, bxDim*byDim*sizeof(float), cudaMemcpyHostToDevice);
	err = cudaMemcpy(C_GPU, _C, cxDim*cyDim*sizeof(float), cudaMemcpyHostToDevice);

	if (err == cudaSuccess) { printf("\nNo problems copying arrays to device!\n");}	

	MatrixMultiplySharedMemory<<<dimGrid, dimBlock, dimBlock.x*dimBlock.x*sizeof(float)>>>(A_GPU, B_GPU, C_GPU, ayDim, cyDim);

	err = cudaMemcpy(_C, C_GPU, axDim*byDim*sizeof(float), cudaMemcpyDeviceToHost);

	if (err == cudaSuccess) { printf("\nNo problems copying arrays from device!\n");}

	cudaFree(A_GPU);
	cudaFree(B_GPU);
	cudaFree(C_GPU);
	cudaDeviceReset();		
	
	printf("\nA\n");
	for(i = 0; i < ax; i++)
	{		
		for(j = 0; j < ay; j++)
		{		
			printf(" %.5f ", A[ay*i+j]);	
		}
		printf("\n");
	}
	
	printf("B\n");
	for(i = 0; i < bx; i++)
	{		
		for(j = 0; j < by; j++)
		{		
			printf(" %.5f ", B[by*i+j]);	
		}
		printf("\n");
	}
	
	printf("C\n");
	for(i = 0; i < cx; i++)
	{		
		for(j = 0; j < cy; j++)
		{		
			printf(" %.5f ", C[cy*i+j]);	
		}
		printf("\n");
	}

	printf("\n_A\n");
	for(i = 0; i < axDim; i++)
	{		
		for(j = 0; j < ayDim; j++)
		{		
			printf(" %.5f ", _A[ayDim*i+j]);	
		}
		printf("\n");
	}
	
	printf("_B\n");
	for(i = 0; i < bxDim; i++)
	{		
		for(j = 0; j < byDim; j++)
		{		
			printf(" %.5f ", _B[byDim*i+j]);	
		}
		printf("\n");
	}
	
	printf("_C\n");
	for(i = 0; i < cxDim; i++)
	{		
		for(j = 0; j < cyDim; j++)
		{		
			printf(" %.5f ", _C[cyDim*i+j]);	
		}
		printf("\n");
	}

	//printf("\nWhere the fuck is the seg fault????");
	return 0;
}
